#pragma once

#include <string>
#include <vector>
#include <iostream>
#include <sstream>
#include <unistd.h>
#include <pwd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <cstdlib>
#include <cstring>

enum RedirectionType
{
    NoneRedir = -1,
    StdinRedir = 0,
    StdoutRedir = 1,
    AppendRedir = 2
};

class Shell
{
public:
    Shell();
    void run();

private:
    std::string getPrompt();
    std::string readCommand();
    void parseCommand(std::string &input);
    bool handleBuiltin();
    void executeCommand();
    void checkRedirection(std::string &input);

    std::vector<std::string> args;
    RedirectionType redirType;
    std::string redirFile;
    int lastStatus;
};

Shell::Shell() : redirType(NoneRedir), lastStatus(0) {}

void Shell::run()
{
    while (true)
    {
        std::string input = readCommand();
        if (input.empty())
            continue;

        parseCommand(input);

        if (args.empty())
            continue;

        if (handleBuiltin())
            continue;

        executeCommand();
    }
}

std::string Shell::getPrompt()
{
    char hostname[256];
    gethostname(hostname, sizeof(hostname));
    std::string user = getlogin();
    char cwd[1024];
    getcwd(cwd, sizeof(cwd));
    return user + "@" + hostname + ":" + cwd + "$ ";
}

std::string Shell::readCommand()
{
    std::cout << getPrompt();
    std::string input;
    std::getline(std::cin, input);
    return input;
}

void Shell::checkRedirection(std::string &input)
{
    redirType = NoneRedir;
    redirFile.clear();

    size_t pos;
    if ((pos = input.find(">>")) != std::string::npos)
    {
        redirType = AppendRedir;
        redirFile = input.substr(pos + 2);
    }
    else if ((pos = input.find(">")) != std::string::npos)
    {
        redirType = StdoutRedir;
        redirFile = input.substr(pos + 1);
    }
    else if ((pos = input.find("<")) != std::string::npos)
    {
        redirType = StdinRedir;
        redirFile = input.substr(pos + 1);
    }

    if (redirType != NoneRedir)
    {
        input = input.substr(0, pos);
        redirFile.erase(0, redirFile.find_first_not_of(" ")); // trim left
    }
}

void Shell::parseCommand(std::string &input)
{
    checkRedirection(input);
    std::istringstream iss(input);
    std::string token;
    args.clear();
    while (iss >> token)
    {
        args.push_back(token);
    }

    if (!args.empty() && args[0] == "ls")
    {
        args.push_back("--color");
    }
}

bool Shell::handleBuiltin()
{
    if (args.empty())
        return false;

    if (args[0] == "cd")
    {
        const char *path = args.size() > 1 ? args[1].c_str() : getenv("HOME");
        chdir(path);
        setenv("PWD", getcwd(nullptr, 0), 1);
        return true;
    }
    else if (args[0] == "export")
    {
        if (args.size() > 1)
        {
            putenv(strdup(args[1].c_str()));
        }
        return true;
    }
    else if (args[0] == "echo")
    {
        if (args.size() == 1)
        {
            std::cout << std::endl;
        }
        else if (args[1][0] == '$')
        {
            if (args[1] == "$?")
            {
                std::cout << lastStatus << std::endl;
                lastStatus = 0;
            }
            else
            {
                const char *val = getenv(args[1].substr(1).c_str());
                if (val)
                    std::cout << val << std::endl;
            }
        }
        else
        {
            for (size_t i = 1; i < args.size(); ++i)
            {
                std::cout << args[i] << " ";
            }
            std::cout << std::endl;
        }
        return true;
    }

    return false;
}

void Shell::executeCommand()
{
    pid_t pid = fork();
    if (pid == 0)
    {
        if (redirType != NoneRedir)
        {
            int fd;
            if (redirType == StdinRedir)
            {
                fd = open(redirFile.c_str(), O_RDONLY);
                dup2(fd, STDIN_FILENO);
            }
            else if (redirType == StdoutRedir)
            {
                fd = open(redirFile.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666);
                dup2(fd, STDOUT_FILENO);
            }
            else if (redirType == AppendRedir)
            {
                fd = open(redirFile.c_str(), O_WRONLY | O_CREAT | O_APPEND, 0666);
                dup2(fd, STDOUT_FILENO);
            }
        }

        std::vector<char *> argv;
        for (auto &arg : args)
            argv.push_back(const_cast<char *>(arg.c_str()));
        argv.push_back(nullptr);

        execvp(argv[0], argv.data());
        perror("execvp failed");
        exit(1);
    }
    else if (pid > 0)
    {
        int status;
        waitpid(pid, &status, 0);
        lastStatus = WEXITSTATUS(status);
    }
}
